import cv2
import numpy as np
from PIL import Image
from picamera2 import Picamera2

from ultralytics import YOLO

# Initialize the Picamera2
picam2 = Picamera2()
picam2.preview_configuration.main.size = (640, 480)
picam2.preview_configuration.main.format = "RGB888"
picam2.preview_configuration.align()
picam2.configure("preview")
picam2.start()

# Load the YOLO11 model
model = YOLO("/ultralytics/yolo11n_ncnn_model")
# model = YOLO("/ultralytics/yolo11n.pt")

print("模型加载完毕，开始识别...")

while True:
    # Capture frame-by-frame
    frame = picam2.capture_array()
    print("获得图像")

    # Run YOLO11 inference on the frame
    results = model(frame)
    print("完成识别")

    # Visualize the results on the frame
    annotated_frame = results[0].plot()
    print("完成标示")

    # image = Image.fromarray(annotated_frame)
    # image.save("detect.jpg")
    # print("完成图像保存")

    # 转换为OpenCV格式（BGR）
    cv2_frame = cv2.cvtColor(np.array(annotated_frame), cv2.COLOR_RGB2BGR)
    print("完成图像格式转换")

    # Display the resulting frame
    # cv2.imshow("Camera With YOLO", cv2_frame)
    # print("完成展示")

    # Break the loop if 'q' is pressed
    if cv2.waitKey(1) == ord("q"):
        break

# Release resources and close windows
cv2.destroyAllWindows()